import os
import torch
from torch.utils.data.dataset import Dataset
import cv2 as cv
from utils.utils import letter_box, image_pretreat, image_translation


class SimpleDatasetClass(Dataset):
    def __init__(self):
        super().__init__()
        self.img_prepath1 = 'E:/dataset/famousface'
        self.img_prepath0 = 'E:/dataset/face/face_train/0'
        self.img_name1 = os.listdir(self.img_prepath1)
        self.img_name0 = os.listdir(self.img_prepath0)

    def __len__(self):
        return len(self.img_name1) + len(self.img_name0)

    def __getitem__(self, item):
        if item < len(self.img_name1):
            img_path = self.img_prepath1 + '/' + self.img_name1[item]
        else:
            img_path = self.img_prepath0 + '/' + self.img_name0[item - len(self.img_name1)]
        image = cv.imread(img_path)
        image = image_pretreat(image, 224)
        label = 1 if item < len(self.img_name1) else 0
        return torch.tensor(image, dtype=torch.float32), torch.tensor(label, dtype=torch.float32)


class SimpleDatasetDetect(Dataset):
    def __init__(self):
        super().__init__()
        self.img_path = 'E:/dataset/face2.png'
        self.x1 = 66
        self.y1 = 54
        self.x2 = 137
        self.y2 = 154
        self.w = self.x2 - self.x1
        self.h = self.y2 - self.y1
        self.cntx = 224 - self.w
        self.cnty = 224 - self.h

    def __len__(self):
        return self.cntx * self.cnty

    def __getitem__(self, item):
        image = cv.imread(self.img_path)
        image = letter_box(image, 224)
        move_y = int(item/self.cntx)
        move_x = item - move_y * self.cntx
        image = image_translation(image, move_x-self.x1, move_y-self.y1)
        image = image[:, :, ::-1].transpose(2, 0, 1)
        image = torch.tensor(image/256, dtype=torch.float32)
        x = move_x + self.w/2
        y = move_y + self.h/2
        label = torch.tensor([x, y, self.w, self.h], dtype=torch.float32)/224
        return image, label
